from rbt.dmu import DecisionMakingUnit


class BnPositionPnlDMU(DecisionMakingUnit):
    """PositionPnlDMU for bn. 
    TODO - add trading fee and impact cost.
    """
    version = "v0"

    def __init__(self):
        super().__init__()
        self.positions = {}  # 用于存储之前的持仓和成本

    def make_decision(self, new_data, previous_result: dict = {}) -> dict:
        pnl_results = {}
        for key, current_position in previous_result.items():
            if key.endswith("_position"):
                position_pnl_key = key.replace("_position", "_pnl")
                if key not in self.positions:
                    self.positions[key] = {"position": 0, "cashflow": 0}

                previous_position = self.positions[key]["position"]
                mid = new_data["close"]  # use cloes price as trading price

                delta_position = current_position - previous_position
                trading_price = mid
                closing_price = mid

                self.positions[key]["position"] = current_position
                prev_cashflow = self.positions[key]["cashflow"]
                total_cashflow = prev_cashflow - delta_position * trading_price
                self.positions[key]["cashflow"] = total_cashflow

                pnl_results[position_pnl_key] = (
                    total_cashflow + closing_price * current_position
                )

        return pnl_results
